﻿using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

public class ChatRoomTools
{
    private ConcurrentDictionary<WebSocket, string> _clients = new ConcurrentDictionary<WebSocket, string>();
    private string _roomId;

    public ChatRoomTools(string roomId)
    {
        _roomId = roomId;
    }

    public async Task AddClient(WebSocket webSocket, string username)
    {
        _clients.TryAdd(webSocket, username);
        //await BroadcastMessage($"{username} 加入了聊天室");

        // 开始接收消息
        _ = ReceiveMessages(webSocket);
    }

    private async Task ReceiveMessages(WebSocket webSocket)
    {
        var buffer = new byte[1024 * 4];

        try
        {
            while (webSocket.State == WebSocketState.Open)
            {
                var result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);

                if (result.MessageType == WebSocketMessageType.Close)
                {
                    await RemoveClient(webSocket);
                }
                else
                {
                    string message = Encoding.UTF8.GetString(buffer, 0, result.Count);
                    string username = _clients[webSocket];
                    await BroadcastMessage($"{username}: {message}");
                }
            }
        }
        catch
        {
            await RemoveClient(webSocket);
        }
    }

    private async Task RemoveClient(WebSocket webSocket)
    {
        if (_clients.TryRemove(webSocket, out string username))
        {
            await webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Connection closed", CancellationToken.None);
            //await BroadcastMessage($"{username} 离开了聊天室");
        }
    }

    private async Task BroadcastMessage(string message)
    {
        byte[] messageBytes = Encoding.UTF8.GetBytes(message);
        var tasks = new List<Task>();

        foreach (var client in _clients)
        {
            if (client.Key.State == WebSocketState.Open)
            {
                tasks.Add(client.Key.SendAsync(
                    new ArraySegment<byte>(messageBytes),
                    WebSocketMessageType.Text,
                    true,
                    CancellationToken.None));
            }
        }

        await Task.WhenAll(tasks);
    }
}